import base64
import requests
import json
import time

class OllamaVisionClient:
    def __init__(self, base_url="http://127.0.0.1:11434", model="qwen3-vl:4b"):
        self.base_url = base_url
        self.model = model
    
    def image_to_base64(self, image_path):
        """将图片转换为 Base64 编码"""
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    
    def describe_image(self, image_path, prompt="请详细描述这张图片的内容", stream=False):
        """
        描述图片内容
        
        Args:
            image_path (str): 图片路径
            prompt (str): 提示词
            stream (bool): 是否使用流式输出
        
        Returns:
            str or generator: 模型响应
        """
        try:
            base64_image = self.image_to_base64(image_path)
            
            data = {
                "model": self.model,
                "prompt": prompt,
                "images": [base64_image],
                "stream": stream
            }
            
            if stream:
                return self._stream_response(data)
            else:
                return self._get_response(data)
                
        except Exception as e:
            return f"错误: {str(e)}"
    
    def _get_response(self, data):
        """获取完整响应"""
        url = f"{self.base_url}/api/generate"
        print("url => ", url)
        response = requests.post(
            url,
            json=data,
            timeout=120
        )
        
        if response.status_code == 200:
            result = response.json()
            return result.get('response', '无响应')
        else:
            return f"请求失败: {response.status_code}"
    
    def _stream_response(self, data):
        """流式响应生成器"""
        response = requests.post(
            f"{self.base_url}/api/generate",
            json=data,
            stream=True,
            timeout=120
        )
        
        if response.status_code == 200:
            for line in response.iter_lines():
                if line:
                    try:
                        json_data = json.loads(line)
                        if 'response' in json_data:
                            yield json_data['response']
                        if json_data.get('done', False):
                            break
                    except json.JSONDecodeError:
                        continue
        else:
            yield f"请求失败: {response.status_code}"
    
    def chat_with_image(self, image_path, initial_question=None):
        """
        与图片进行多轮对话
        """
        base64_image = self.image_to_base64(image_path)
        conversation_history = []
        
        print("开始与图片对话（输入 'quit' 退出）")
        print("=" * 50)
        
        while True:
            if initial_question and not conversation_history:
                user_input = initial_question
                initial_question = None
            else:
                user_input = input("\n你的问题: ").strip()
            
            if user_input.lower() in ['quit', 'exit', '退出']:
                break
            if not user_input:
                continue
            
            # 构建对话上下文
            messages = []
            for role, content in conversation_history:
                messages.append({"role": role, "content": content})
            
            # 添加当前问题（始终包含图片）
            data = {
                "model": self.model,
                "prompt": user_input,
                "images": [base64_image],
                "stream": False
            }
            
            print("AI: ", end="", flush=True)
            response = self._get_response(data)
            print(response)
            
            # 保存对话历史
            conversation_history.append(("user", user_input))
            conversation_history.append(("assistant", response))

# 使用示例
if __name__ == "__main__":
    client = OllamaVisionClient()
    
    # 示例1：基础描述
    image_path = "/Users/hhwang/Downloads/3e0.jpg"  # 替换为你的图片路径
    
    print("=== 基础图片描述 ===")
    description = client.describe_image(image_path)
    print(description)
    
    print("\n" + "="*60)
    
    # # 示例2：提出具体问题
    # print("=== 具体问题回答 ===")
    # question = "图片中的主要颜色是什么？有哪些物体？"
    # answer = client.describe_image(image_path, question)
    # print(f"问题: {question}")
    # print(f"回答: {answer}")
    
    # print("\n" + "="*60)
    
    # # 示例3：流式输出
    # print("=== 流式输出示例 ===")
    # print("描述: ", end="", flush=True)
    # for chunk in client.describe_image(image_path, stream=True):
    #     print(chunk, end="", flush=True)
    # print()  # 换行
    
    # print("\n" + "="*60)
    
    # 示例4：多轮对话（取消注释以使用）
    # print("=== 多轮对话模式 ===")
    # client.chat_with_image(image_path, "请先描述一下这张图片")
